# distutils: language = c
# cython: language_level=3

import numpy as np
cimport numpy as np
from libc.stdlib cimport malloc, free

cdef extern from "c_ldtw.h":
    ctypedef struct PathNode:
        int i, j
    ctypedef struct LDTWResult:
        PathNode *path
        int path_len
        double dist

    LDTWResult ldtw_path(const double *s1, size_t len1,
                         const double *s2, size_t len2,
                         int max_length)

    LDTWResult ldtw_path_tw(const double *s1, size_t len1,
                         const double *s2, size_t len2,
                         int max_length,
                         double beta,
                         const double *gap_max_vec);

def ldtw_path_c(np.ndarray[np.float64_t, ndim=1] s1,
                                    np.ndarray[np.float64_t, ndim=1] s2,
                                    int max_length = -1):
        # 若用户没传，用最大长度作为默认值
    if max_length == -1:
        max_length = max(s1.size, s2.size)
    min_len = max(s1.size, s2.size)
    if max_length < min_len:
        raise ValueError(f"max_length({max_length}) must be >= max(len(s1), len(s2)) = {min_len}")

    cdef LDTWResult res
    res = ldtw_path(&s1[0], s1.size,
                    &s2[0], s2.size,
                    max_length)
    if res.path_len == 0 or res.dist < 0:
        raise ValueError("C ldtw_path failed: probably max_length too small.")

    try:
        path_list = [(res.path[k].i, res.path[k].j)
                     for k in range(res.path_len)]
    finally:
        free(res.path)

    return path_list, res.dist

# 时间加权LDTW
# 逐段可变惩罚版本
def ldtw_path_tw_c(np.ndarray[np.float64_t, ndim=1] s1,
                                             np.ndarray[np.float64_t, ndim=1] s2,
                                             int max_length = -1,
                                             double beta=1.0,
                                             gap_max=None):
    if gap_max is None:
        gap_max = np.full(s1.size, max(s1.size, s2.size), dtype=np.float64)
    else:
        gap_max = np.asarray(gap_max, dtype=np.float64)
        if gap_max.ndim == 0:                 # 标量
            gap_max = np.full(s1.size, gap_max, dtype=np.float64)
        elif gap_max.ndim != 1 or gap_max.size != s1.size:
            raise ValueError("gap_max 长度必须等于 s1 的长度")
    cdef np.ndarray[np.float64_t, ndim=1] gap_max_vec = gap_max

    # 若用户没传，用最大长度作为默认值
    if max_length == -1:
        max_length = max(s1.size, s2.size)
    min_len = max(s1.size, s2.size)
    if max_length < min_len:
        raise ValueError(f"max_length({max_length}) must be >= max(len(s1), len(s2)) = {min_len}")
    
    cdef LDTWResult res
    res = ldtw_path_tw(&s1[0], s1.size,
                           &s2[0], s2.size,
                           max_length, beta,
                           &gap_max_vec[0])
    if res.path_len == 0 or res.dist < 0:
        raise ValueError("C ldtw_path_tw failed: probably max_length too small.")

    try:
        path_list = [(res.path[k].i, res.path[k].j)
                     for k in range(res.path_len)]
    finally:
        free(res.path)

    return path_list, res.dist